///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/wfn/wfn.cpp
///
/// \brief   implementation file for the wavefunction class
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<complex>
#include<iostream>
#include<string>

#include<formic/wfn/wfn.h>
#include<formic/exception.h>
#include<formic/archive.h>
#include<formic/output/output.h>
#include<formic/mpi/interface.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Checks that a name is not already used among the wavefunction's constituents.
///
/// \param[in]      name     the name to check
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void formic::Wavefunction<S>::check_new_name(const std::string & name) const {
  if ( _wi_map.count(name) > 0 || _wf_map.count(name) > 0 || _co_map.count(name) > 0 )
    throw formic::Exception("the wavefunction already contains an object with name \"%s\"") % name;
}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Adds a coefficient object to the wavefunction.
///
/// \param[in]      c        a pointer to the coefficient object to add
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void formic::Wavefunction<S>::add_coeff_obj(const boost::shared_ptr<formic::CoeffObj<S> > & c) {

  // make sure the name is not already in use
  this->check_new_name(c->name());

  // add the coefficient object to the name-coefficient mapping
  _co_map[c->name()] = c;

  // report what we added
  if (formic::mpi::rank() == 0)
    formic::of << boost::format("Added a coefficient object of type %s to the wavefunction under the name \"%s\"")
                    % c->class_name()
                    % c->name()
               << std::endl
               << std::endl;

}

template void formic::Wavefunction<double>::add_coeff_obj(const boost::shared_ptr<formic::CoeffObj<double> > &);
template void formic::Wavefunction<std::complex<double> >::add_coeff_obj(const boost::shared_ptr<formic::CoeffObj<std::complex<double> > > &);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Adds a data object to the wavefunction
///
/// \param[in,out]  d             a pointer to the data object to add
/// \param[in,out]  next_fix_pos  the position in which to place the next set of fixed data
/// \param[in,out]  next_var_pos  the position in which to place the next set of variable data
/// \param[in,out]  m             the mapping between names and data objects
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> template <class T>
void formic::Wavefunction<S>::add_wfn_data(const boost::shared_ptr<formic::WfnData<T> > & d,
                                           int & next_fix_pos,
                                           int & next_var_pos,
                                           std::map<std::string, boost::shared_ptr<formic::WfnData<T> > > & m) {

  // make sure the name is not already in use
  this->check_new_name(d->name());

//  // set the data position
//  if (d->fixed()) {
//    d->set_pos(next_fix_pos);
//    next_fix_pos -= d->size();
//  } else {
//    d->set_pos(next_var_pos);
//    next_var_pos += d->size();
//  }

  // add the data object to the name-data mapping
  assert( m.count(d->name()) == 0 );
  m[d->name()] = d;

  // print a message stating that the data was added
  struct XX {
    static std::string get(int x) { return "integer"; }
    static std::string get(double x) { return "real"; }
    static std::string get(std::complex<double>  x) { return "complex"; }
  };
  if (formic::mpi::rank() == 0)
    formic::of << boost::format("Added %s wavefunction data under the name \"%s\"") % XX::get(T()) % d->name() << std::endl << std::endl;

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Reads a data object from the supplied archive and adds it to the wavefunction
///
/// \param[in,out]  arch          an archive from which to read the data object
/// \param[in,out]  next_fix_pos  the position in which to place the next set of fixed data
/// \param[in,out]  next_var_pos  the position in which to place the next set of variable data
/// \param[in,out]  m             the mapping between names and data objects
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> template <class T>
void formic::Wavefunction<S>::add_wfn_data(formic::Archive & arch,
                                           int & next_fix_pos,
                                           int & next_var_pos,
                                           std::map<std::string, boost::shared_ptr<formic::WfnData<T> > > & m) {

  // create the data object
  boost::shared_ptr<formic::WfnData<T> > d( new formic::WfnData<T>(arch) );

  // add it to the wavefunction
  this->add_wfn_data(d, next_fix_pos, next_var_pos, m);
  
}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Adds an integer data object to the wavefunction.
///
/// \param[in,out]  arch     an archive from which to read the data object
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void formic::Wavefunction<S>::add_wfn_int_data(formic::Archive & arch) {
  this->add_wfn_data(arch, _fixed_int_pos, _var_int_pos, _wi_map);
}

template void formic::Wavefunction<double>::add_wfn_int_data(formic::Archive &);
template void formic::Wavefunction<std::complex<double> >::add_wfn_int_data(formic::Archive &);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Adds a floating point data object to the wavefunction.
///
/// \param[in,out]  arch     an archive from which to read the data object
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::Wavefunction<S>::add_wfn_float_data(formic::Archive & arch) {
  this->add_wfn_data(arch, _fixed_float_pos, _var_float_pos, _wf_map);
}

template void formic::Wavefunction<double>::add_wfn_float_data(formic::Archive &);
template void formic::Wavefunction<std::complex<double> >::add_wfn_float_data(formic::Archive &);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Adds a floating point data object to the wavefunction.
///
/// \param[in,out]  d        a pointer to the data object to add
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::Wavefunction<S>::add_wfn_data(const boost::shared_ptr<formic::WfnData<S> > & d) {
  this->add_wfn_data(d, _fixed_float_pos, _var_float_pos, _wf_map);
}

template void formic::Wavefunction<double>::add_wfn_data(const boost::shared_ptr<formic::WfnData<double> > &);
template void formic::Wavefunction<std::complex<double> >::add_wfn_data(const boost::shared_ptr<formic::WfnData<std::complex<double> > > &);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Adds an integer data object to the wavefunction.
///
/// \param[in,out]  d        a pointer to the data object to add
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::Wavefunction<S>::add_wfn_data(const boost::shared_ptr<formic::WfnData<int> > & d) {
  this->add_wfn_data(d, _fixed_int_pos, _var_int_pos, _wi_map);
}

template void formic::Wavefunction<double>::add_wfn_data(const boost::shared_ptr<formic::WfnData<int> > &);
template void formic::Wavefunction<std::complex<double> >::add_wfn_data(const boost::shared_ptr<formic::WfnData<int> > &);

// ///////////////////////////////////////////////////////////////////////////////////////////////////
// /// \brief  returns a reference to the named floating point wavefunction data
// ///
// /// \param[in]      name     the name of the data we want
// ///
// ///////////////////////////////////////////////////////////////////////////////////////////////////
// template <class S> formic::WfnData<S> & formic::Wavefunction<S>::get_fdata(const std::string & name) {
//   if ( _wf_map.count(name) == 0 )
//     throw formic::Exception("no floating point data object with the name \"%s\" in Wavefunction::get_fdata.") % name;
//   return *_wf_map[name];
// }
// 
// template formic::WfnData<double> & formic::Wavefunction<double>::get_fdata(const std::string &);
// template formic::WfnData<std::complex<double> > & formic::Wavefunction<std::complex<double> >::get_fdata(const std::string &);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  returns a shared pointer to the named coefficient object
///
/// \param[in]      name     the name of the coefficient object we want
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> boost::shared_ptr<formic::CoeffObj<S> > formic::Wavefunction<S>::get_co(const std::string & name) {
  if ( _co_map.count(name) == 0 )
    throw formic::Exception("no coefficient object with the name \"%s\" in Wavefunction::get_co.") % name;
  return _co_map[name];
}

template boost::shared_ptr<formic::CoeffObj<double> > formic::Wavefunction<double>::get_co(const std::string &);
template boost::shared_ptr<formic::CoeffObj<std::complex<double> > > formic::Wavefunction<std::complex<double> >::get_co(const std::string &);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  returns a shared pointer to the named floating point data object
///
/// \param[in]      name     the name of the data object we want
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> boost::shared_ptr<formic::WfnData<S> > formic::Wavefunction<S>::get_wf(const std::string & name) {
  if ( _wf_map.count(name) == 0 )
    throw formic::Exception("no floating point data object with the name \"%s\" in Wavefunction::get_wf.") % name;
  return _wf_map[name];
}

template boost::shared_ptr<formic::WfnData<double> > formic::Wavefunction<double>::get_wf(const std::string &);
template boost::shared_ptr<formic::WfnData<std::complex<double> > > formic::Wavefunction<std::complex<double> >::get_wf(const std::string &);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  returns a shared pointer to the named integer data object
///
/// \param[in]      name     the name of the data object we want
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> boost::shared_ptr<formic::WfnData<int> > formic::Wavefunction<S>::get_wi(const std::string & name) {
  if ( _wi_map.count(name) == 0 )
    throw formic::Exception("no integer data object with the name \"%s\" in Wavefunction::get_wi.") % name;
  return _wi_map[name];
}

template boost::shared_ptr<formic::WfnData<int> > formic::Wavefunction<double>::get_wi(const std::string &);
template boost::shared_ptr<formic::WfnData<int> > formic::Wavefunction<std::complex<double> >::get_wi(const std::string &);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  writes the wavefunction's data objects to a binary archive
///
/// \param[in]      name     the name of the data object we want
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template <class S> void formic::Wavefunction<S>::dump_data_to_bin_arch(const std::string & filename) {

  // open the binary archive
  formic::BinaryFileArchive arch(filename, "wt");

  // write integer data
  for (typename IMap::iterator it = _wi_map.begin(); it != _wi_map.end(); it++)
    it->second->to_command(arch);

  // write floating point data
  for (typename FMap::iterator it = _wf_map.begin(); it != _wf_map.end(); it++)
    it->second->to_command(arch);

}

template void formic::Wavefunction<double>::dump_data_to_bin_arch(const std::string &);
template void formic::Wavefunction<std::complex<double> >::dump_data_to_bin_arch(const std::string &);
